from flask import Flask, request, jsonify, render_template
import torch
import torch.nn as nn
from torchvision import transforms
from PIL import Image
import io
import json
import os


# 定义AlexNet模型类
class AlexNet(nn.Module):
    def __init__(self, num_classes=8, dropout_rate=0.5):
        super(AlexNet, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 96, kernel_size=11, stride=4, padding=2),
            nn.ReLU(inplace=True),
            nn.LocalResponseNorm(size=5, alpha=0.0001, beta=0.75, k=2),
            nn.MaxPool2d(kernel_size=3, stride=2),

            nn.Conv2d(96, 256, kernel_size=5, padding=2),
            nn.ReLU(inplace=True),
            nn.LocalResponseNorm(size=5, alpha=0.0001, beta=0.75, k=2),
            nn.MaxPool2d(kernel_size=3, stride=2),

            nn.Conv2d(256, 384, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),

            nn.Conv2d(384, 384, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),

            nn.Conv2d(384, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
        )
        self.avgpool = nn.AdaptiveAvgPool2d((6, 6))
        self.classifier = nn.Sequential(
            nn.Dropout(p=dropout_rate),
            nn.Linear(256 * 6 * 6, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(p=dropout_rate),
            nn.Linear(4096, 4096),
            nn.ReLU(inplace=True),
            nn.Linear(4096, num_classes),
        )

    def forward(self, x):
        x = self.features(x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)
        return x


app = Flask(__name__)

# 全局变量
model = None
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
class_names = ['飞机', '车', '猫', '狗',
               '花', '水果', '摩托', '人']

# 图像预处理
transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])


def load_model(model_path):
    """加载训练好的模型"""
    global model
    try:
        model = AlexNet(num_classes=8)

        # 加载模型权重
        checkpoint = torch.load(model_path, map_location=device)

        if 'model_state_dict' in checkpoint:
            model.load_state_dict(checkpoint['model_state_dict'])
        elif 'state_dict' in checkpoint:
            model.load_state_dict(checkpoint['state_dict'])
        else:
            model.load_state_dict(checkpoint)

        model.to(device)
        model.eval()
        print(f"模型加载成功，使用设备: {device}")
        return True
    except Exception as e:
        print(f"模型加载失败: {str(e)}")
        return False


def predict_image(image):
    """对单张图像进行预测"""
    try:
        # 预处理图像
        image_tensor = transform(image).unsqueeze(0).to(device)

        # 预测
        with torch.no_grad():
            outputs = model(image_tensor)
            probabilities = torch.nn.functional.softmax(outputs, dim=1)
            confidence, predicted = torch.max(probabilities, 1)

        # 返回结果
        result = {
            'class_id': predicted.item(),
            'class_name': class_names[predicted.item()],
            'confidence': confidence.item(),
            'all_probabilities': {
                class_names[i]: probabilities[0][i].item()
                for i in range(len(class_names))
            }
        }
        return result
    except Exception as e:
        print(f"预测错误: {str(e)}")
        return None


@app.route('/')
def index():
    """主页"""
    return render_template('index.html')


@app.route('/predict', methods=['POST'])
def predict():
    """预测接口"""
    try:
        # 检查是否有文件上传
        if 'file' not in request.files:
            return jsonify({'error': '没有上传文件'}), 400

        file = request.files['file']

        # 检查文件名
        if file.filename == '':
            return jsonify({'error': '没有选择文件'}), 400

        # 检查文件类型
        if not file.content_type.startswith('image/'):
            return jsonify({'error': '请上传图像文件'}), 400

        # 读取图像
        image_bytes = file.read()
        image = Image.open(io.BytesIO(image_bytes)).convert('RGB')

        # 预测
        result = predict_image(image)

        if result is None:
            return jsonify({'error': '预测失败'}), 500

        return jsonify(result)

    except Exception as e:
        return jsonify({'error': f'服务器错误: {str(e)}'}), 500


@app.route('/health', methods=['GET'])
def health_check():
    """健康检查接口"""
    return jsonify({'status': 'healthy', 'device': str(device)})


if __name__ == '__main__':
    # 加载模型
    model_path = './alexnet_model.pth'
    if not load_model(model_path):
        print("无法加载模型，请检查模型文件路径")
        exit(1)

    # 启动Flask应用
    app.run(host='0.0.0.0', port=5000, debug=False)